%%%%%%%%%%%%%%%%%%%%%%%
%%
%% Module summary
%%
%% Test library
%%
%% Contains useful functions when
%% writing tests.
%%
%%%%%%%%%%%%%%%%%%%%%%%

-module(test_lib).
-export([get_random_db_node/0]).
-export([socket_start/0,socket_start/1,socket_stop/0,
         socket_put/2,socket_put_list/1,
         socket_get/1,socket_get_list/1]).
-export([socket_start_pid/0,socket_start_pid/1,socket_stop_pid/1,
         socket_put_pid/3,socket_put_list_pid/2,
         socket_get_pid/2,socket_get_list_pid/2]).
-export([rpc_put/3,rpc_put_list/2,rpc_get/2,rpc_get_list/2]).
-export([random_key_value_pairs/1,random_key_value_pairs/2,
         random_string/1]).
-export([all_db_nodes/0]).


-define(DEFAULT_STRING_LENGTH,10).


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% General node handling functions
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% Returns a list of all connected 
% erlang nodes, including itself
all_nodes() ->
    [node()] ++ nodes().

% Returns a list of all nodes currently 
% running a database instance
all_db_nodes() ->
    [Node || Node <- all_nodes(), is_running_db(Node)].

% Returns a random database node
get_random_db_node() ->
    <<A:32, B:32, C:32>> = crypto:strong_rand_bytes(12),
    random:seed(A,B,C),
    DB_Nodes = all_db_nodes(),
    case length(DB_Nodes) > 0 of
        true  -> lists:nth(random:uniform(length(DB_Nodes)), DB_Nodes);
        false -> {error,not_running}
    end.

% Check if node is running the database application,
% returns boolean()
is_running_db(Node) ->
    case rpc:call(Node,application,which_applications,[]) of
        {badrpc,_}   -> false;
        Running_apps -> lists:keymember(db,1,Running_apps)
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% Socket connections
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Start a socket connection to a random db node
socket_start() ->
    case get_random_db_node() of
        {error,_} -> io:fwrite("test_lib: Could not find any DB node\n");
        DB_Node -> Port = get_tcp_port(DB_Node),
                   Address = get_node_address(DB_Node),
                   clientovich:start(Address,Port)
    end.

% Establish a TCP connection to Node
socket_start(Node) ->
    Address = get_node_address(Node),
    case get_tcp_port(Node) of
        {error, _} -> io:fwrite("test_lib: Could not connect to " ++ Address 
                                                                  ++ "\n");
        Port       -> clientovich:start(Address,Port)
    end.

% Stop client (closes socket aswell)    
socket_stop() ->
    clientovich:stop().

% Put a key and value using the socket connection
socket_put(Key, Value) ->
    clientovich:put(Key,Value).

% Put a list of {Key, Value} pair into DB
socket_put_list(List) ->
    [clientovich:put(Key,Value) || {Key,Value} <- List].

% Get a key
socket_get(Key) ->
    clientovich:get(Key).

% Query the database using a list of {Key,Value} pairs
socket_get_list(List) ->
    [clientovich:get(Key) || {Key,_} <- List].

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% Socket connections, using clientovich2
%%% 
%%% Used for multiple slaves on the same node
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Start a socket connection to a random db node
socket_start_pid() ->
    case get_random_db_node() of
        {error,_} -> io:fwrite("test_lib: Could not find any DB node\n");
        DB_Node -> Port = get_tcp_port(DB_Node),
                   Address = get_node_address(DB_Node),
                   io:format("Client connecting to: ~w~n",[DB_Node]),
                   clientovich2:start(Address,Port)
    end.

% Establish a TCP connection to Node
socket_start_pid(Node) ->
    Address = get_node_address(Node),
    case get_tcp_port(Node) of
        {error, _} -> io:fwrite("test_lib: Could not connect to " ++ Address 
                                                                  ++ "\n");
        Port       -> clientovich2:start(Address,Port)
    end.

% Stop client (closes socket aswell)    
socket_stop_pid(Pid) ->
    clientovich2:stop(Pid).

% Put a key and value using the socket connection
socket_put_pid(Pid,Key, Value) ->
    clientovich2:put(Pid,Key,Value).

% Put a list of {Key, Value} pair into DB
socket_put_list_pid(Pid,List) ->
    [clientovich2:put(Pid,Key,Value) || {Key,Value} <- List].

% Get a key
socket_get_pid(Pid,Key) ->
    clientovich2:get(Pid,Key).

% Query the database using a list of {Key,Value} pairs
socket_get_list_pid(Pid,List) ->
    [clientovich2:get(Pid,Key) || {Key,_} <- List].
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% Socket helper functions
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Find the current tcp port used by Node
get_tcp_port(Node) ->
    case is_running_db(Node) of
        true  ->
            {ok,Port} = rpc:call(Node, application, get_env, [db,port]),
            Port;
        false ->
            {error, db_not_running}
    end.

% Parse the address from nodename
get_node_address(Node) ->
    String = atom_to_list(Node),
    [_Name,Address] = string:tokens(String, "@"),
    Address. 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% RPC queries
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Insert a key-value pair into db using rpc
rpc_put(Node, Key, Value) ->
    case rpc:call(Node, node, put, [Key, Value]) of
        {badrpc,_} -> io:fwrite("test_lib: Node is unreachable");
        Reply      -> Reply
    end.
    

% Insert a list of key-value pairs of type {Key,Value}
rpc_put_list(Node, List) ->
    [rpc:call(Node, node, put, [Key, Value]) || {Key,Value} <- List].

% Get a value with given key from db, using rpc call
rpc_get(Node, Key) ->
    case rpc:call(Node, node, get, [Key]) of
        {badrpc,_} -> io:fwrite("test_lib: Node is unreachable");
        Reply      -> Reply
    end.

% Query the database using a list of {Key,Value} pairs
rpc_get_list(Node, List) ->
    [rpc:call(Node, node, get, [Key]) || {Key,_} <- List].



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% Random string generation
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Generate a list of random strings
random_key_value_pairs(Amount) ->
    [{random_string(?DEFAULT_STRING_LENGTH),
      random_string(?DEFAULT_STRING_LENGTH)} || _ <- lists:seq(1,Amount)].

% Generate a list of random string of wanted length
random_key_value_pairs(Amount, Length) ->
    [{random_string(Length),random_string(Length)}
    || _ <- lists:seq(1,Amount)].

% Returns a random string of given length
random_string(Length) ->
    % Generate a seed
    <<A:32, B:32, C:32>> = crypto:strong_rand_bytes(12),
    random_string_help(Length,all_characters(),{A,B,C}).

% Recursive helper function for random_string
random_string_help(0,_,_) -> [];
random_string_help(Length,AllowedChars,Seed) ->
   {NewInt,NewSeed} = random:uniform_s(length(AllowedChars),Seed),
   [lists:nth(NewInt,AllowedChars)] ++ 
            random_string_help(Length-1,AllowedChars,NewSeed).

% Returns a string of characters
all_characters() ->
    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWQYZ".    

